from manim import *
from pathlib import Path
import os

FLAGS = f"-qh"
SCENE = "Scene"

Text.set_default(font_size = 28)
BraceLabel.set_default(brace_config={"fill_opacity": 0.6, "sharpness": 0.8})

Y_OFFSET = 2.5
X_OFFSET = 4
COLORS = [BLUE, GREEN, RED, YELLOW]

class Scene(Scene):
    def create_letter_code_objects(self, letters, codes, y_offset=Y_OFFSET, x_offset=X_OFFSET):
        letter_objects = []
        code_objects = []
        for i, letter in enumerate(letters):
            letter_obj = Text(letter).shift(LEFT * x_offset + UP * (y_offset - i))
            code_obj = Text(codes[letter]).shift(LEFT * (x_offset-2) + UP * (y_offset - i))
            letter_objects.append(letter_obj)
            code_objects.append(code_obj)
        return letter_objects, code_objects

    def create_brace_labels(self, letter_obj, code_obj, text_letter, text_code):
        br_letter = BraceLabel(letter_obj, text=text_letter).next_to(letter_obj, DOWN)
        br_code = BraceLabel(code_obj, text=text_code).next_to(code_obj, DOWN)
        br_letter.label.set_color(COLORS[1])
        br_code.label.set_color(COLORS[0])
        return br_letter, br_code

    def construct(self):

        # Data for three scenes
        data = [
            ({"A": "000", "B": "001", "C": "010", "D": "011", "E": "100"}, ["A", "B", "C", "D", "E"], "1", "3", "5", r"3/1=3"),
            ({"AA": "00000", "AB": "00001", "AC": "00010", "...": "...", "EE": "11000"}, ["AA", "AB", "AC", "...", "EE"], "2", "5", "25", r"5/2=2.5"),
            ({"AAA": "0000000", "AAB": "0000001", "AAC": "0000010", "...": "...", "EEE": "1111100"}, ["AAA", "AAB", "AAC", "...", "EEE"], "3", "7", "125", r"7/3 \approx 2.33"),
        ]

        # Entropy text
        entropy_text = Text("熵最大值: ").shift(UP * Y_OFFSET + RIGHT * 1)
        entropy_num  = MathTex(r"\log_2(5) \approx 2.32").next_to(entropy_text, RIGHT * 2)
        entropy_num[0][5].set_color(COLORS[3])
        entropy_num[0][8:12].set_color(COLORS[2])

        self.wait(20)

        prev_letter_objects = prev_code_objects = None
        prev_br_letter = prev_br_code = prev_br = None

        for i, (dict_, letters, text_letter, text_code, brace_text, equation) in enumerate(data):
            letter_objects, code_objects = self.create_letter_code_objects(letters, dict_)
            br_letter, br_code = self.create_brace_labels(letter_objects[-1], code_objects[-1], text_letter, text_code)
            square = Square(side_length=4.0)
            br = BraceLabel(square, brace_direction=LEFT, buff=0.6, text=brace_text)
            if brace_text == "5":
                br.label.set_color(COLORS[3])
            br.shift(LEFT*(X_OFFSET-2) + UP * (Y_OFFSET-2))
            
            if i == 0:
                # Animation1
                single_text = Text("方式一: ").next_to(entropy_text, DOWN * 2)
                single_num = MathTex(equation).next_to(single_text, RIGHT * 2)
                single_num[0][0].set_color(COLORS[0])
                single_num[0][2].set_color(COLORS[1])
                single_num[0][4].set_color(COLORS[2])
                
                self.wait(20)
                self.play(*[FadeIn(letter) for letter in letter_objects], run_time=1.5)
                self.wait(15)
                self.play(*[Write(code) for code in code_objects], run_time=1.5)
                self.wait(10)
                self.play(FadeIn(br_letter), FadeIn(br_code), run_time=1.5)
                self.wait(10)
                self.play(FadeIn(single_text), TransformMatchingTex(Group(br_letter.label, br_code.label).copy(), single_num), run_time=1.5)
                self.wait(20)
                self.play(FadeIn(br))
                self.wait(20)
                self.play(FadeIn(entropy_text), TransformMatchingTex(br.label.copy(), entropy_num), run_time=1.5)
                self.wait(20)
            else:
                # Replace previous objects with new ones
                self.play(
                    ReplacementTransform(prev_br_letter, br_letter),
                    *[ReplacementTransform(prev_letter, curr_letter) for prev_letter, curr_letter in zip(prev_letter_objects, letter_objects)],
                    run_time=1.5
                )
                self.wait(20)
                self.play(ReplacementTransform(prev_br.label, br.label))
                self.wait(20)
                self.play(
                    ReplacementTransform(prev_br_code, br_code),
                    *[ReplacementTransform(prev_code, curr_code) for prev_code, curr_code in zip(prev_code_objects, code_objects)],
                    run_time=1.5
                )
                self.wait(20)

                if i == 1:
                    double_text = Text("方式二: ").next_to(single_text, DOWN * 2)
                    double_num = MathTex(equation).next_to(double_text, RIGHT * 2)
                    double_num[0][0].set_color(COLORS[0])
                    double_num[0][2].set_color(COLORS[1])
                    double_num[0][4:7].set_color(COLORS[2])
                    self.play(FadeIn(double_text), TransformMatchingTex(Group(br_letter.label, br_code.label).copy(), double_num), run_time=1.5)
                else:
                    triple_text = Text("方式三: ").next_to(double_text, DOWN * 2)
                    triple_num = MathTex(equation).next_to(triple_text, RIGHT * 2)
                    triple_num[0][0].set_color(COLORS[0])
                    triple_num[0][2].set_color(COLORS[1])
                    triple_num[0][4:8].set_color(COLORS[2])
                    self.play(FadeIn(triple_text), TransformMatchingTex(Group(br_letter.label, br_code.label).copy(), triple_num), run_time=1.5)
                    self.wait(20)
                    self.play(
                        Indicate(entropy_num[0][8:12], scale_factor=1.4 , color=COLORS[2]),
                        Indicate(single_num[0][4]    , scale_factor=1.4 , color=COLORS[2]),
                        Indicate(double_num[0][4:7]  , scale_factor=1.4 , color=COLORS[2]),
                        Indicate(triple_num[0][4:8]  , scale_factor=1.4 , color=COLORS[2]),
                        run_time=2
                    )
                self.wait(20)
            
            prev_letter_objects = letter_objects
            prev_code_objects = code_objects
            prev_br_letter = br_letter
            prev_br_code = br_code
            prev_br = br

        # Final animation wait
        self.wait(20)
        self.play(
            *[FadeOut(mob)for mob in self.mobjects],
            run_time=2
        )
        self.wait(20)


        ########## END ##########
        chart = BarChart(
            values=[0.25, 0.25, 0.25, 0.25, 0.25],
            bar_names=["A", "B", "C", "D", "E"],
            y_range=[0, 1, 0.2],
            y_length=4,
            x_length=6,
            x_axis_config={"font_size": 36},
        )

        c_bar_lbls = chart.get_bar_labels(font_size=48)


        chart_2 = BarChart(
            values=[0.25, 0.5, 0.125, 0.10, 0.025],
            bar_names=["A", "B", "C", "D", "E"],
            y_range=[0, 1, 0.2],
            y_length=4,
            x_length=6,
            x_axis_config={"font_size": 36}
        )

        c2_bar_lbls = chart_2.get_bar_labels(font_size=48)

        self.wait(20)
        self.play(
            FadeIn(chart),
            FadeIn(c_bar_lbls),
            run_time=2
        )
        self.wait(20)
        self.play(
            ReplacementTransform(chart, chart_2),
            ReplacementTransform(c_bar_lbls, c2_bar_lbls),
            run_time = 2.5
        )
        self.wait(20)












if __name__  == '__main__':
    script_name = f"{Path(__file__).resolve()}"
    os.system(f"manim {FLAGS} {script_name} {SCENE}")